Skip to content

fix(loss): min/max packed extrema instead of summing them - #3901

Open
tianyi-zhang-02 wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
tianyi-zhang-02:fix/packed-metric-extrema
Open

fix(loss): min/max packed extrema instead of summing them#3901
tianyi-zhang-02 wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
tianyi-zhang-02:fix/packed-metric-extrema

Conversation

@tianyi-zhang-02

@tianyi-zhang-02 tianyi-zhang-02 commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes packed min/max diagnostics that were being summed as additive metrics.

SequencePackingLossWrapper and the worker aggregators recognize extrema by the explicit _min / _max suffix contract. Names that merely contain those substrings in the middle remain additive. Current main's registered MOPD exception, opd_full_decomposition_error, keeps its max reduction even though it does not use the suffix. Fully masked value samples retain the existing sentinel behavior, and an all-sentinel critic step maps to -1.0 like policy metrics.

Losses and gradients are unchanged.

Validation

Current head 4a5af89017bd2bf51a9fa4cd32ba81ec59de7f20, rebased onto upstream main at 90a2a212d503455d8590be5c8de3cb989d3425b0.

  • packed metric and critic aggregation regressions: 8 passed
  • covers suffix extrema, embedded-substring additive names, the registered MOPD exception, and all-sentinel behavior
  • Ruff check, Ruff format check, and git diff --check: passed

An earlier pre-refresh head also passed the gated Megatron split-parity test on two real H100s. It ran two optimizer steps through freshly initialized sync and split policies and compared the loss curve, grad norm, and every reduced per-microbatch metric.

topology result elapsed
DP=2, TP=1 1 passed 494.22 s
DP=1, TP=2 1 passed 367.62 s

Environment: 2× NVIDIA H100 80 GB, nvcr.io/nvidia/nemo-rl:v0.7.0, Python 3.13, PyTorch 2.11.0+cu130, CUDA 13, and NCCL_NVLS_ENABLE=0. The current-head rebase and MOPD regression were verified locally; the 2-GPU parity test was not rerun after that refresh.

@tianyi-zhang-02
tianyi-zhang-02 requested review from a team as code owners August 28, 2026 22:28
@copy-pr-bot

copy-pr-bot Bot commented Aug 28, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@tianyi-zhang-02
tianyi-zhang-02 requested review from a team as code owners August 30, 2026 19:57
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Aug 31, 2026
SequencePackingLossWrapper folds per-sequence metric dicts into one. It
special-cases extrema through a hardcoded allowlist naming only the four
probs_ratio keys, so MseValueLossFn's values_min/values_max fall through
to '+=' and are summed. Three packed sequences spanning -3..9 report
values_min=4.0 -- a positive number for a critic whose predictions go
negative -- and values_max=15.0. The error grows with packing density,
so it is not a stable offset a reader could correct for.

The loss function already says what these are: 'Min/max are per-MB;
ppo.py takes min/max across MBs.' Five other sites apply the
'_min'/'_max' suffix rule to this very dict, and one of them --
megatron_value_worker.py:611 -- is this wrapper's own direct consumer,
skipping the divide because it is an extremum while the wrapper upstream
has already summed it. Use the same rule here.

That alone is not enough. MseValueLossFn returns 0.0 for a fully-masked
sequence where ClippedPGLossFn returns +/-inf, and 0.0 is a plausible
value that wins the min against an all-positive critic: the one-line
version reports 0.0 where the truth is 3.0. sample_mask is
loss_multiplier, which overlong_filtering zeroes per sample, and under
packing one filtered sample in a pack is enough. So the sentinel moves
to +/-inf and ppo.py skips it, matching what that file already does for
probs_ratio at :1774 and :2757.

Metrics only -- packed and unpacked losses are bit-identical. Reachable
on ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-pp2cp2-pack, which
nightly.txt runs.

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
@tianyi-zhang-02
tianyi-zhang-02 force-pushed the fix/packed-metric-extrema branch from 43f85ae to 4a5af89 Compare September 12, 2026 18:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants